package com.imooc.springcloud.rules;

import com.netflix.client.config.IClientConfig;
import com.netflix.loadbalancer.AbstractLoadBalancerRule;
import com.netflix.loadbalancer.IRule;
import com.netflix.loadbalancer.Server;
import lombok.NoArgsConstructor;
import org.springframework.util.CollectionUtils;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.io.UnsupportedEncodingException;
import java.security.MessageDigest;
import java.security.NoSuchAlgorithmException;
import java.util.List;
import java.util.SortedMap;
import java.util.TreeMap;

/**
 * @author 小木
 * @date 2021/3/22
 * @desc 自定义Rule
 */
@NoArgsConstructor
public class MyRule extends AbstractLoadBalancerRule implements IRule {

    @Override
    public void initWithNiwsConfig(IClientConfig clientConfig) {

    }

    @Override
    public Server choose(Object key) {
        HttpServletRequest request = ((ServletRequestAttributes) RequestContextHolder.getRequestAttributes()).getRequest();
        // 这里可以根据具体业务选择合适的【针对性】，比如：根据用户id【强绑定】的话，可以添加用户id
        String uri = request.getServletPath() + "/" + request.getQueryString();
        return route(uri.hashCode(), getLoadBalancer().getAllServers());
    }

    /**
     * 选择合适的【服务节点】
     */
    public Server route(int hashId, List<Server> allServers) {
        Server server = null;
        if (CollectionUtils.isEmpty(allServers)) {
            return null;
        }
        TreeMap<Long, Server> serverTreeMap = new TreeMap<>();
        // 虚化若干个节点放到环上，注意：此处【非均 匀分布】
        allServers.forEach(item -> {
            for (int i = 0; i < 8; i++) {
                long hash = hash(item.getId() + i);
                serverTreeMap.put(hash, item);
            }
        });
        // 再次进行hash => long
        long hash = hash(String.valueOf(hashId));

        // 根据【key】获取【所有>该key】的元素
        SortedMap<Long, Server> serverSortedMap = serverTreeMap.tailMap(hash);
        // 注意：为空的情况可能是因为 request URL 生成的hash > 任意一个服务器对应的hash值（也就是环尾），那么【取环首】
        if (serverSortedMap.isEmpty()) {
            server = serverTreeMap.firstEntry().getValue();
        } else {
            server = serverSortedMap.get(serverSortedMap.firstKey());
        }
        return server;
    }

    /**
     * 根据服务器id，计算hash值
     */
    public long hash(String key) {
        MessageDigest md5;
        try {
            md5 = MessageDigest.getInstance("MD5");
        } catch (NoSuchAlgorithmException e) {
            throw new RuntimeException(e);
        }
        byte[] bytes = null;
        try {
            bytes = key.getBytes("UTF-8");
        } catch (UnsupportedEncodingException e) {
            throw new RuntimeException(e);
        }
        md5.update(bytes);
        // 生成摘要【16位的byte】
        byte[] digest = md5.digest();
        return (digest[2] & 0xff << 16) | (digest[1] & 0xff << 8) | (digest[0] & 0xff);
    }

}
